#include "npdm_patterns.h"
#include "npdm_operators.h"
#include "npdm_sweep.h"
#include "npdm_tester.h"

using namespace Npdm;

//===========================================================================================================================================================
Npdm_tester::Npdm_tester(int order, int nsites, Npdm_sweep npdm_sweep)
:
order_(order),
nsites_(nsites),
npdm_sweep_(npdm_sweep)
{
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------
// Doesn't take into account spin-permutation redundancy

long int Npdm_tester::get_3pdm_spin_degeneracy( std::vector<int> & indices )
{

  long int valid_spin = 0;

  std::set< std::tuple<CD,int,int> > opset;
  for (int i=0; i<2; i++) {
    for (int j=0; j<2; j++) {
      for (int k=0; k<2; k++) {
        for (int l=0; l<2; l++) {
          for (int m=0; m<2; m++) {
            for (int n=0; n<2; n++) {

              // Need same number of alpha and beta creation and destruction 
              if ( i+j+k != l+m+n ) continue;
    
              opset.clear();
              opset.insert( std::make_tuple(CREATION,    indices.at(0), i) );
              opset.insert( std::make_tuple(CREATION,    indices.at(1), j) );
              opset.insert( std::make_tuple(CREATION,    indices.at(2), k) );
              opset.insert( std::make_tuple(DESTRUCTION, indices.at(3), l) );
              opset.insert( std::make_tuple(DESTRUCTION, indices.at(4), m) );
              opset.insert( std::make_tuple(DESTRUCTION, indices.at(5), n) );
              if ( opset.size() == 6 ) valid_spin++;
    
            }
          }
        }
      }
    }
  }
//  std::cout << "-------------\n";
//std::cout << valid_spin << std::endl;

  assert ( valid_spin <= 20 );
  return valid_spin;
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------
// Doesn't take into account spin-permutation redundancy

long int Npdm_tester::get_2pdm_spin_degeneracy( std::vector<int> & indices )
{

  long int valid_spin = 0;

  std::set< std::tuple<CD,int,int> > opset;
  for (int i=0; i<2; i++) {
    for (int j=0; j<2; j++) {
      for (int k=0; k<2; k++) {
        for (int l=0; l<2; l++) {

          // Need same number of alpha and beta creation and destruction 
          if ( i + j != k + l ) continue;

//          std::cout << i << j << k << l << std::endl;
          opset.clear();
          opset.insert( std::make_tuple(CREATION,    indices.at(0), i) );
          opset.insert( std::make_tuple(CREATION,    indices.at(1), j) );
          opset.insert( std::make_tuple(DESTRUCTION, indices.at(2), k) );
          opset.insert( std::make_tuple(DESTRUCTION, indices.at(3), l) );
          if ( opset.size() == 4 ) {
            valid_spin++;
//print it
//            for ( auto it = opset.begin(); it != opset.end(); it++ ) {
//              std::cout << "(" << std::get<0>(*it) << std::get<1>(*it) << std::get<2>(*it) << ")";
//            }
//            std::cout << "\n";
          }

        }
      }
    }
  }
//  std::cout << "-------------\n";
//std::cout << valid_spin << std::endl;

  assert ( valid_spin <= 6 );
  return valid_spin;
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

bool Npdm_tester::valid_fermion_string( std::vector<int> & indices )
{
  // Assume indices are sorted and first half creation and second half destruction

  // Check if no more than two creation with same index
  int k = 1;
  for (int i=1; i < indices.size()/2; i++) {
    if ( indices[i] == indices[i-1] )
      k++;
    else
      k=1;
    if ( k > 2 ) return false;
  }

  // Check if no more than two destruction with same index
  k = 1;
  for (int i=1+indices.size()/2; i < indices.size(); i++) {
    if ( indices[i] == indices[i-1] )
      k++;
    else
      k=1;
    if ( k > 2 ) return false;
  }

  return true;
}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

std::pair<long int, long int> Npdm_tester::test_if_4pdm_fully_built()
{

  long int count = 0;
  for (int i=1; i<=nsites_; i++) {
    for (int j=i; j<=nsites_; j++) {
      for (int k=j; k<=nsites_; k++) {
        for (int l=k; l<=nsites_; l++) {
          for (int m=1; m<=nsites_; m++) {
            for (int n=m; n<=nsites_; n++) {
              for (int p=n; p<=nsites_; p++) {
                for (int q=p; q<=nsites_; q++) {
                  if ( std::make_tuple(i,j,k,l) > std::make_tuple(m,n,p,q) ) continue;
                  std::vector<int> indices = {i,j,k,l,m,n,p,q};
                  if ( not valid_fermion_string( indices ) ) continue;
                  count++;
                  // Only print if not built in DMRG algorithm
                  if ( npdm_sweep_.included_in_dmrg_sweep( indices ) ) continue;
                  std::cout << "Not built! ";
                  std::cout << i << j << k << l << m << n << p << q << std::endl;
                  assert ( false );
                }
              }
            }
          }
        }
      }
    }
  }
  return std::make_pair(count,0);

}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

std::pair<long int, long int> Npdm_tester::test_if_3pdm_fully_built()
{

  long int count = 0;
  long int spin = 0;
  for (int i=1; i<=nsites_; i++) {
    for (int j=i; j<=nsites_; j++) {
      for (int k=j; k<=nsites_; k++) {
        for (int l=1; l<=nsites_; l++) {
          for (int m=l; m<=nsites_; m++) {
            for (int n=m; n<=nsites_; n++) {
              if ( std::make_tuple(i,j,k) > std::make_tuple(l,m,n) ) continue;
              std::vector<int> indices = {i,j,k,l,m,n};
              if ( not valid_fermion_string( indices ) ) continue;
              count++;
              spin = spin + get_3pdm_spin_degeneracy( indices );
              // Only print if not built in DMRG algorithm
              if ( npdm_sweep_.included_in_dmrg_sweep( indices ) ) continue;
              std::cout << "Not built! ";
              std::cout << i << j << k << l << m << n << std::endl;
              assert ( false );
            }
          }
        }
      }
    }
  }
  return std::make_pair(count,spin);

}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

std::pair<long int, long int> Npdm_tester::test_if_2pdm_fully_built()
{

  long int count = 0;
  long int spin = 0;
  for (int i=1; i<=nsites_; i++) {
    for (int j=i; j<=nsites_; j++) {
      for (int k=1; k<=nsites_; k++) {
        for (int l=k; l<=nsites_; l++) {
          if ( std::make_tuple(i,j) > std::make_tuple(k,l) ) continue;
          std::vector<int> indices = {i,j,k,l};
          if ( not valid_fermion_string( indices ) ) continue;
          count++;
          spin = spin + get_2pdm_spin_degeneracy( indices );
          // Only print if not built in DMRG algorithm
          if ( npdm_sweep_.included_in_dmrg_sweep( indices ) ) continue;
          std::cout << "Not built! ";
          std::cout << i << j << k << l << std::endl;
          assert ( false );
        }
      }
    }
  }
  return std::make_pair(count,spin);

}

//-----------------------------------------------------------------------------------------------------------------------------------------------------------

void Npdm_tester::test_if_fully_built( int order ) 
{

  std::pair<long int, long int> nstrings;

  if ( order == 2 )
    nstrings = test_if_2pdm_fully_built();
  else if ( order == 3 )
    nstrings = test_if_3pdm_fully_built();
  else if ( order == 4 )
    nstrings = test_if_4pdm_fully_built();
  else
    assert (false);

  std::cout << "Number of irreducible spatial operator strings = "<< nstrings.first << std::endl;;
  std::cout << "Number of spin-orbital operator strings        = "<< nstrings.second << std::endl;;

}

//===========================================================================================================================================================

